package com.springboot.demo.lock.zookeeper;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import lombok.extern.log4j.Log4j2;
import org.I0Itec.zkclient.IZkDataListener;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkNoNodeException;

/**
 * 基础分布式锁, 使用zookeeper实现的分布式锁
 *
 * @author supu
 * @date 2019-03-21 9:42
 **/
@Log4j2
public class BaseDistributedLock {
    /**
     * zk 客户端
     */
    private final ZkClient zkClient;
    /**
     * 用于保存Zookeeper中实现分布式锁的节点，例如/locker节点，该节点是个持久节点，在该节点下面创建临时顺序节点来实现分布式锁
     */
    private final String basePath;
    /**
     * 同basePath一样
     */
    private final String path;
    /**
     * 锁名称前缀，/locker下创建的顺序节点，例如以lock-开头，这样便于过滤无关节点
     */
    private final String lockName;
    /**
     * 最大重试次数
     */
    private static final int MAX_RETRY_COUNT = 10;

    public BaseDistributedLock(ZkClient zkClient, String path, String lockName) {
        this.zkClient = zkClient;
        this.basePath = path;
        this.path = path.concat("/").concat(lockName);
        this.lockName = lockName;
        if(!zkClient.exists(basePath)){
            zkClient.createPersistent(basePath);
        }
    }

    /**
     * 超时时间内获取锁
     *
     * @param startMillis 开始时间
     * @param timeToWait  超时时间
     * @param path        节点路径
     * @return
     * @throws Exception
     */
    private boolean waitToLock(long startMillis, Long timeToWait, String path) throws Exception {
        // 获取锁标识
        boolean hasLock = false;
        // 删除锁标识
        boolean doDelete = false;

        try {
            while (!hasLock) {
                log.warn("{} 开始获取锁了", path);
                // 获取path下的所有从小到大排好序的顺序节点
                List<String> children = getSortedChildren();
                // 获取子节点名称, 例如 /demo/lock_0002 子节点名称是lock_0002
                String sequenceNodeName = path.substring(basePath.length() + 1);
                //
                int nodeIndex = children.indexOf(sequenceNodeName);

                /*
                 * 如果在getSortedChildren中没有找到之前创建的[临时]顺序节点，这表示可能由于网络闪断而导致
                 * Zookeeper认为连接断开而删除了我们创建的节点，此时需要抛出异常，让上一级去处理
                 * 上一级的做法是捕获该异常，并且执行重试指定的次数，见后面的 attemptLock方法
                 */
                if (nodeIndex < 0) {
                    throw new ZkNoNodeException("节点没有找到: " + sequenceNodeName);
                }

                // 序号最小说明获取到了锁
                if (nodeIndex == 0) {
                    hasLock = true;
                    break;
                }

                // 没有获取到锁，监听比自己次小的那个节点, 例如 /demo/lock_0001
                String pathToWatch = children.get(nodeIndex - 1);
                // 如果次节点被删除了，则表示当前节点就是最小的了，所以使用 countDownLatch来等待获取锁
                String previousSequencePath = basePath.concat("/").concat(pathToWatch);
                final CountDownLatch countDownLatch = new CountDownLatch(1);
                final IZkDataListener previousListener = new IZkDataListener() {
                    /**
                     * 监听指定节点删除时触发该方法
                     */
                    @Override
                    public void handleDataDeleted(String dataPath)
                            throws Exception {
                        // 次小节点删除事件发生时，让countDownLatch结束等待
                        // 此时还需要重新让程序回到while，重新判断一次！
                        countDownLatch.countDown();
                    }

                    /**
                     * 监听指定节点的数据发生变化触发该方法
                     *
                     */
                    @Override
                    public void handleDataChange(String dataPath,
                                                 Object data) throws Exception {

                    }

                };

                try {
                    zkClient.subscribeDataChanges(previousSequencePath, previousListener);
                    if (timeToWait != null) {
                        timeToWait -= (System.currentTimeMillis() - startMillis);
                        startMillis = System.currentTimeMillis();
                        if (timeToWait <= 0) {
                            // 已经超时
                            doDelete = true;
                            break;
                        }
                        countDownLatch.await(timeToWait, TimeUnit.MILLISECONDS);
                    } else {
                        countDownLatch.await();
                    }
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    zkClient.unsubscribeDataChanges(previousSequencePath, previousListener);
                }
            }
        } catch (ZkNoNodeException e) {
            // 发生异常，需要删除节点
            doDelete = true;
            throw e;
        } finally {
            // 如果需要删除该节点，则删除该节点
            if (doDelete) {
                deletePath(path);
            }
        }

        return hasLock;
    }

    /**
     * 获取锁
     *
     * @param time     超时时间
     * @param timeUnit 时间单位
     * @return
     * @throws Exception
     */
    protected String attemptLock(long time, TimeUnit timeUnit) throws Exception {
        long startMills = System.currentTimeMillis();
        Long timeToWait = timeUnit != null ? timeUnit.toMillis(time) : null;

        // 创建的临时节点path
        String nodePath = null;
        // 是否获取到了锁
        boolean hasLock = false;
        // 是否完成了获取锁
        boolean isDone = false;
        // 重试次数
        int retryCount = 0;

        while (!isDone) {
            isDone = true;

            try {
                // 创建临时有序节点
                nodePath = createEphemeralSequential();

                /**
                 * 该方法用于判断自己是否获取到了锁，即自己创建的顺序节点在lock的所有子节点中是否最小
                 * 如果没有获取到锁，则等待其它客户端锁的释放，并且稍后重试直到获取到锁或者超时
                 */
                hasLock = waitToLock(startMills, timeToWait, nodePath);
            } catch (Exception e) {
                if (retryCount++ < MAX_RETRY_COUNT) {
                    isDone = false;
                } else {
                    throw e;
                }
            }
        }

        log.warn("{} 获取锁 {}", nodePath, (hasLock ? "成功" : "失败"));

        if (hasLock) {
            return nodePath;
        }

        return null;
    }

    /**
     * 释放锁
     *
     * @param lockPath 锁路径
     * @throws Exception
     */
    protected void releaseLock(String lockPath) throws Exception {
        deletePath(lockPath);
    }

    /**
     * 创建临时有序节点
     *
     * @return
     * @throws Exception
     */
    private String createEphemeralSequential() throws Exception {
        return zkClient.createEphemeralSequential(path, null);
    }

    /**
     * 删除节点
     *
     * @param path
     * @throws Exception
     */
    private void deletePath(String path) throws Exception {
        zkClient.delete(path);
    }

    /**
     * 获取排序后的所有子节点
     *
     * @return
     */
    private List<String> getSortedChildren() {
        try {
            List<String> children = zkClient.getChildren(basePath);
            Collections.sort(children, (o1, o2) -> {
                return getLockNumber(o1, lockName).compareTo(getLockNumber(o2, lockName));
            });

            return children;
        } catch (Exception e) {
            zkClient.createPersistent(basePath, true);
            return getSortedChildren();
        }
    }

    /**
     * 获取锁序号
     *
     * @param path     节点路径
     * @param lockName 锁前缀
     * @return
     */
    private String getLockNumber(String path, String lockName) {
        int index = path.lastIndexOf(lockName);
        if (index >= 0) {
            index += lockName.length();
            return index <= path.length() ? path.substring(index) : "";
        }
        return path;
    }
}
